#ifndef _KEYDISTRIBUTIONCENTER_H_
#define _KEYDISTRIBUTIONCENTER_H_

#include <cstring>
#include <map>
#include "Blowfish.h"
#include "PracticalSocket.h"

// keyS|idA|, keyLength + idLength < 30. // 32
#define INNER_BUFFER_LENGTH 64
// enc[keyS|idA|idB|nonceA|enc[keyS|idA|]|], idA|idB|nonceA, keyLength * 2 + idLength * 3 + nonceLength < 12 < 121. // 128
#define MAX_BUFFER_LENGTH 128

/*
 * Comaprator for maps.
 */
struct cmp_str {
	bool operator() (char const *a, char const *b) {
		return strcmp(a, b) < 0;
	}
};

/*
 * Class which represents a server.
 */
class Server {
	private:
		char* serverNumber;
		char* keyB;
	public:
		Server(char*, char*);
		~Server();
		char* getServerNumber();
		char* getKeyB();
};
/*
 * Server Constructor.
 */
Server::Server(char* serverNumberIn, char* keyBIn) {
	serverNumber = serverNumberIn;
	keyB = keyBIn;
}
/*
 * Server Destructor.
 */
Server::~Server() {
	delete[] serverNumber;
	delete[] keyB;
}
/*
 * Server getServerNumber.
 */
char* Server::getServerNumber(){
	return serverNumber;
}
/*
 * Server getKeyB.
 */
char* Server::getKeyB() {
	return keyB;
}

/*
 * Class which represents a client and its connections.
 */
class Client {
	private:
		map<char*, char*, cmp_str> connections; // Map of server numbers and keyS.
		char* keyA;
	public:
		Client(char*);
		~Client();
		void addConnection(char*, char*);
		char* getConnectionKey(char*);
		char* getKeyA();
};
/*
 * Client Constructor.
 */
Client::Client(char* keyAIn) {
	keyA = keyAIn;
}
/*
 * Client Destructor.
 */
Client::~Client() {
	delete &connections;
	delete[] keyA;
}
/*
 * Client addConnection.
 */
void Client::addConnection(char* serverNumber, char* keyS) {
	connections.insert(pair<char*, char*>(serverNumber, keyS));
}
/*
 * Client getConnectionKey.
 */
char* Client::getConnectionKey(char* serverNumber) {
	map<char*, char*, cmp_str>::iterator iter = connections.find(serverNumber);
	if (iter != connections.end()) {
		return iter->second;
	} else {
		cout << serverNumber << " is not in connections map" << endl;
		return 0;
	}
}
/*
 * Client getKeyA.
 */
char* Client::getKeyA() {
	return keyA;
}

/*
 * Class which represents the KeyDistributionCenter and it's maps.
 */
class KeyDistributionCenter {
	private:
		map<char*, Server*, cmp_str> idServerMap; // Map of server id's and servers.
		map<char*, Client*, cmp_str> idClientMap; // Map of client id's and clients.
		pthread_mutex_t parseMutex; // For strtok
		int encryptionOn; // 1 = true.
		Server* getServer(char*);
		Client* getClient(char*);
		void parseRequest(char*, char**, char**, char**);
		char* buildResponse(char*, char*, char*, char*, char*, char*);
	public:
		KeyDistributionCenter(int);
		~KeyDistributionCenter();
		void addServer(char*, Server*);
		void addClient(char*, Client*);
		void handleTCPRequest(TCPSocket*);
};
/*
 * KeyDistributionCenter getServer.
 */
Server* KeyDistributionCenter::getServer(char* idB) {
	map<char*, Server*, cmp_str>::iterator iter = idServerMap.find(idB);
	if (iter != idServerMap.end()) {
		return iter->second;
	} else {
		cout << idB << " is not in server map" << endl;
		return NULL;
	}
}
/*
 * KeyDistributionCenter getClient.
 */
Client* KeyDistributionCenter::getClient(char* idA) {
	map<char*, Client*, cmp_str>::iterator iter = idClientMap.find(idA);
	if (iter != idClientMap.end()) {
		return iter->second;
	} else {
		cout << idA << " is not in client map" << endl;
		return NULL;
	}
}
/*
 * KeyDistributionCenter parseRequest.
 * Parse message "idA|idB|nonceA|".
 * Not thread safe.
 */
void KeyDistributionCenter::parseRequest(char* requestBuffer, char** idA, char** idB, char** nonceA) {
	sprintf(*idA, "%s", strtok(requestBuffer, "|"));
	sprintf(*idB, "%s", strtok(NULL, "|"));
	sprintf(*nonceA, "%s", strtok(NULL, "|"));
}
/*
 * KeyDistributionCenter buildResponse.
 * Builds response message along with encryption.
 */
char* KeyDistributionCenter::buildResponse(char* idA, char* idB, char* keyA, char* keyB, char* keyS, char* nonceA) {
	// Create message for idB.
	char* innerBuffer = new char[INNER_BUFFER_LENGTH];
	memset(innerBuffer, 0, INNER_BUFFER_LENGTH);
	// keyS|idA|
	sprintf(innerBuffer, "%s|%s|", keyS, idA);
	// If encryption is set to on.
	Blowfish blowfish;
	if (encryptionOn == 1) {
		blowfish.Set_Passwd(keyB);
		blowfish.Encrypt(innerBuffer, (!(INNER_BUFFER_LENGTH % 8 == 0) ? INNER_BUFFER_LENGTH + 8 - (INNER_BUFFER_LENGTH % 8) : INNER_BUFFER_LENGTH));
	}
	// Create response message.
	char* responseBuffer = new char[MAX_BUFFER_LENGTH];
	memset(responseBuffer, 0, MAX_BUFFER_LENGTH);
	// keyS|idA|idB|nonceA|innerBuffer|
	sprintf(responseBuffer, "%s|%s|%s|%s|%s|", keyS, idA, idB, nonceA, innerBuffer);
	// If encryption is set to on.
	if (encryptionOn == 1) {
		blowfish.Set_Passwd(keyA);
		blowfish.Encrypt(responseBuffer, (!(MAX_BUFFER_LENGTH % 8 == 0) ? MAX_BUFFER_LENGTH + 8 - (MAX_BUFFER_LENGTH % 8) : MAX_BUFFER_LENGTH));
	}
	delete[] innerBuffer;
	return responseBuffer;
}
/*
 * KeyDistributionCenter handleTCPRequest.
 * TCP client handling function.
 * Modified from PracticalSocket's TCPEchoServer-Thread.cpp example.
 */
void KeyDistributionCenter::handleTCPRequest(TCPSocket* socket) {
	// Get request from socket.
	char* requestBuffer = new char[MAX_BUFFER_LENGTH];
	memset(requestBuffer, 0, MAX_BUFFER_LENGTH);
	socket->recv(requestBuffer, MAX_BUFFER_LENGTH);
	// Parse request.
	char* idA = new char[INNER_BUFFER_LENGTH];
	memset(idA, 0, INNER_BUFFER_LENGTH);
	char* idB = new char[INNER_BUFFER_LENGTH];
	memset(idB, 0, INNER_BUFFER_LENGTH);
	char* nonceA = new char[INNER_BUFFER_LENGTH];
	memset(nonceA, 0, INNER_BUFFER_LENGTH);
	pthread_mutex_lock(&parseMutex);
	parseRequest(requestBuffer, &idA, &idB, &nonceA);
	pthread_mutex_unlock(&parseMutex);
	// Release requestBuffer resource.
	delete[] requestBuffer;
	// Print request.
	cout << "Recd from " << idA << ":\n\tidA|idB|nonceA|\n\t" << idA << "|" << idB << "|" << nonceA << "|" << endl;
	// Get objects from map.
	Client* client = getClient(idA);
	Server* server = getServer(idB);
	// Get keyA & keyB.
	char* keyA = client->getKeyA();
	char* keyB = server->getKeyB();
	// Get keyS.
	char* keyS = client->getConnectionKey(server->getServerNumber());
	// Build response message.
	char* responseBuffer;
	responseBuffer = buildResponse(idA, idB, keyA, keyB, keyS, nonceA);
	// Send message back to process.
	socket->send(responseBuffer, MAX_BUFFER_LENGTH);
	// Print out information to be sent.
	if (encryptionOn == 1) {
		cout << "Sent to " << idA << ":\n\tE(Ka)[Ks|idA|idB|nonceA|E(Kb)[Ks|idA|]|\n\t" << responseBuffer << endl;
	} else {
		cout << "Sent to " << idA << ":\n\tKs|idA|idB|nonceA|Ks|idA||\n\t" << responseBuffer << endl;
	}
	// Destructor closes socket.
	// Free resources.
	delete[] idA;
	delete[] idB;
	delete[] nonceA;
	delete[] responseBuffer;
}
/*
 * KeyDistributionCenter Constructor.
 */
KeyDistributionCenter::KeyDistributionCenter(int encryptionIn) {
	pthread_mutex_init(&parseMutex, NULL);
	encryptionOn = encryptionIn;
}
/*
 * KeyDistributionCenter Destructor.
 */
KeyDistributionCenter::~KeyDistributionCenter() {
	delete &idServerMap;
	delete &idClientMap;
}
/*
 * KeyDistributionCenter addServer.
 */
void KeyDistributionCenter::addServer(char* idB, Server* server) {
	idServerMap.insert(pair<char*, Server*>(idB, server));
}
/*
 * KeyDistributionCenter addClient.
 */
void KeyDistributionCenter::addClient(char* idA, Client* client) {
	idClientMap.insert(pair<char*, Client*>(idA, client));
}

#endif
